local config = require "config"
local _M = {}
local log = ngx.log
local ERR = ngx.ERR

function open_redis()
    local redis = require "resty.redis"
    local red = redis:new()
    red:set_timeout(1000) -- 1 sec
    local ok, err = red:connect(config["redis_host"], config["redis_port"])
    if not ok then
        log(ERR, "redis failed to connect: ", err)
    end
    if config["redis_password"] ~= '' then
        ok, err = red:auth(config["redis_password"])
        if not ok then
            log(ERR, "redis auth to connect: ", err)
        end
    end
    return red, err
end

function return_redis(red)
    -- put it into the connection pool of size 100,
    -- with 10 seconds max idle time
    local ok, err = red:set_keepalive(10000, 100)
    if not ok then
        log(ERR, "failed to set keepalive: ", err)
        return
    end
end

function _M.limit(api_info)
    local red, err = open_redis()
    if err then
        return
    end

    -- IP黑名单
    local ret = checkIpBlacklist(api_info["ipBlacklist"])
    if ret ~= 200 then
        return_redis(red)
        ngx.exit(ret)
    end

    local uri, uriSeconds, uriTimes, ipUriSeconds, ipUriTimes = api_info["request_uri"], api_info["uri_limit_seconds"], api_info["uri_limit_times"], api_info["ip_uri_limit_seconds"], api_info["ip_uri_limit_times"]
    if uriSeconds > 0 then
        local ret = checkUriLimit(red, uri, uriSeconds, uriTimes)
        if ret ~= 200 then
            return_redis(red)
            ngx.exit(ret)
        end
    end

    if ipUriSeconds > 0 then
        local ret = checkIpUriLimit(red, uri, ipUriSeconds, ipUriTimes)
        if ret ~= 200 then
            return_redis(red)
            ngx.exit(ret)
        end
    end

    return_redis(red)
end

function checkUriLimit(red, uri, uriSeconds, uriTimes)
    -- 针对整个接口限流
    local value, _ = red:get(uri)
    if value == ngx.null then
        red:setex(uri, uriSeconds, 1)
    else
        local count = tonumber(value)
        if count >= uriTimes then
            return 403
        end
        red:incr(uri)
    end
    return 200
end

function get_ip()
    local headers = ngx.req.get_headers()
    local ip = headers["X-REAL-IP"] or headers["X_FORWARDED_FOR"] or ngx.var.remote_addr or "0.0.0.0"
    return ip
end

function checkIpUriLimit(red, uri, ipUriSeconds, ipUriTimes)
    -- 针对IP接口限流
    local ip = get_ip()
    local value, _ = red:get(ip .. uri)
    if value ~= ngx.null then
        local count = tonumber(value)
        if count >= ipUriTimes then
            return 403
        end
        red:incr(ip .. uri)
    else
        red:setex(ip .. uri, ipUriSeconds, 1)
    end
    return 200
end

function checkIpBlacklist(ipBlacklist)
    local ipBlacklist_count = table.getn(ipBlacklist)
    if ipBlacklist_count == 0 then
        return 200
    end
    local ip = get_ip()
    for _, ipblack in pairs(ipBlacklist) do
        if ip == ipblack["ip"] then
            return 403
        end
    end
    return 200
end

return _M